package cn.dengta.webapp.base.util;

import java.util.Map;
import java.util.UUID;
import java.util.concurrent.Callable;

import me.codeplayer.util.StringUtil;
import org.slf4j.*;

public class MDCUtil {

	public static final String TRACE_ID = "traceId";
	public static final String PARENT_ID = "parentId";
	private static final Logger log = LoggerFactory.getLogger(MDCUtil.class);

	public static String getTraceId() {
		String traceId = MDC.get(TRACE_ID);
		return traceId == null ? "" : traceId;
	}

	public static void setTraceId(String traceId) {
		MDC.put(TRACE_ID, traceId);
	}

	public static void setParentId(String traceId) {
		MDC.put(PARENT_ID, traceId);
	}

	public static void remove() {
		MDC.remove(TRACE_ID);
	}

	public static void clear() {
		MDC.clear();
	}

	public static String generateTraceId() {
		return UUID.randomUUID().toString().replace("-", "");
	}

	public static void setTraceId() {
		setTraceId(generateTraceId());
	}

	public static String setTraceIdIfAbsent() {
		String traceId = getTraceId();
		if (StringUtil.isEmpty(traceId)) {
			traceId = generateTraceId();
			setTraceId(traceId);
		}
		return traceId;
	}

	public static String setTraceIdIfAbsent(String parentId) {
		setParentId(parentId);
		return setTraceIdIfAbsent();
	}

	public static <T> Callable<T> wrap(final Callable<T> callable) {
		return wrap(callable, null);
	}

	/**
	 * 用于父线程向线程池中提交任务时，将自身MDC中的数据复制给子线程
	 */
	public static <T> Callable<T> wrap(final Callable<T> callable, final Map<String, String> context) {
		return new Callable<T>() {
			@Override
			public T call() throws Exception {
				if (context == null) {
					MDC.clear();
				} else {
					MDC.setContextMap(context);
				}
				setTraceIdIfAbsent();
				try {
					return callable.call();
				} finally {
					MDC.clear();
				}
			}
		};
	}

	public static Runnable wrap(final Runnable runnable) {
		return wrap(runnable, null);
	}

	public static Runnable wrapParentTraceId(final Runnable runnable) {
		return wrapParentContext(runnable, setTraceIdIfAbsent());
	}

	public static Runnable wrapParentContext(final Runnable runnable, String parentId) {
		return wrap(runnable, Map.of(TRACE_ID, generateTraceId(), PARENT_ID, parentId));
	}

	/**
	 * 用于父线程向线程池中提交任务时，将自身MDC中的数据复制给子线程
	 */
	public static Runnable wrap(final Runnable runnable, final Map<String, String> context) {
		return () -> {
			if (context == null) {
				MDC.clear();
			} else {
				MDC.setContextMap(context);
			}
			setTraceIdIfAbsent();
			try {
				runnable.run();
			} finally {
				log.info("任务执行结束");
				MDC.clear();
			}
		};
	}

}
